from darknet import *
import time
from prune_utils.prune_utils import *
import argparse
from prune_utils.get_prune_cfg import *
from terminaltables import AsciiTable

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--cfg', type=str, default='cfg/yolov2_plate.cfg', help="cfg file path")
    parser.add_argument('--weights', type=str, default='weights/000020.weights', help='sparse model weights')
    parser.add_argument('--percent', type=float, default=0.5, help='channel prune percent')
    parser.add_argument('--img_size', type=int, default=416, help='inference size (pixels)')
    opt = parser.parse_args()
    print(opt)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = Darknet(opt.cfg).to(device)
    img_size=opt.img_size
    if opt.weights.endswith('.pt'):
        model.load_state_dict(torch.load(opt.weights)['model'])
    else:
        model.load_weights(opt.weights)
    print('\n loaded weights from', opt.weights)

    obtain_num_parameters = lambda model: sum([param.nelement() for param in model.parameters()])
    oritin_nparameters = obtain_num_parameters(model)
    CBL_idx, Conv_idx, prune_idx = GetModelPruneDefs(model.blocks)
    thre_value = GetPruneThre(model, prune_idx, opt.percent)
    num_filters, filters_mask = GetFiltersMask(model, thre_value, CBL_idx, prune_idx, 0.3)
    CBLidx2mask={idx:mask.astype('float32') for idx,mask in zip(CBL_idx,filters_mask)}
    #生成新的压缩模型对象
    blocks_del_net=model.blocks.copy()
    blocks_del_net.pop(0)
    compact_module_def=deepcopy(blocks_del_net)
    for idx,num in zip(CBL_idx,num_filters):
        assert compact_module_def[idx]['type']=='convolutional'
        compact_module_def[idx]['filters']=str(num)

    compact_model=Darknet([model.hyperparams.copy()]+compact_module_def,(img_size,img_size)).to(device)
    compact_nparameters=obtain_num_parameters(compact_model)
    InitWeightsFromLooseModel(compact_model,model,CBL_idx,Conv_idx,CBLidx2mask)
    metri_table=[
        ['Metric','Before','After'],
        ['Parameters',f'{oritin_nparameters}',f'{compact_nparameters}']
    ]
    print(AsciiTable(metri_table).table)
    #保存网络配置文件跟权重
    SavePruneRes([model.hyperparams.copy()],compact_model,opt.cfg,opt.weights,opt.percent)






































